#include "scan_patch.h"

// 全局变量定义
PatchEntry patch_table[0x300] = {};
int patch_count = 0;
SpcEntry spc_table[0x500] = {};
int spc_count = 0;
struct x86_instruction_t* x86_it_table;

PatchEntry Restore_Patch_Value_test_patch_entry;

bool search_for_spc(uint32_t addr_) {
    int i = 0;
    int limit = spc_count - (spc_count % 4); // 处理剩余部分

    // 4 次展开循环，减少分支预测开销
    for (; i < limit; i += 4) {
        __builtin_prefetch(&spc_table[i + 4], 0, 1); // 预取数据，提高缓存命中率

        if ((addr_ >= spc_table[i].start_addr   && addr_ <= spc_table[i].end_addr) ||
            (addr_ >= spc_table[i + 1].start_addr && addr_ <= spc_table[i + 1].end_addr) ||
            (addr_ >= spc_table[i + 2].start_addr && addr_ <= spc_table[i + 2].end_addr) ||
            (addr_ >= spc_table[i + 3].start_addr && addr_ <= spc_table[i + 3].end_addr)) {
            return true;
        }
    }

    // 处理剩余的元素
    for (; i < spc_count; i++) {
        if (addr_ >= spc_table[i].start_addr && addr_ <= spc_table[i].end_addr) {
            return true;
        }
    }
    return false;
}


PatchEntry* Restore_Patch_Value (uint32_t addr_) {
    int i = 0;
    for (; i < patch_count; i++) {
        if (addr_ == patch_table[i].patch_addr) {
            return  &(patch_table[i]);
        }
    }
    return &Restore_Patch_Value_test_patch_entry;
}

void x86_scan(void* addr_) {
    struct x86_instruction_t* x86_i_t = NULL;
    PatchEntry* test_patch_entry = NULL;
    uint32_t test_addr;
    bool state = false;
    // if ((uint32_t)addr_ == (uint32_t)(0x00001edc)) {
    //     put_str(" 1edc \n");
    //     asm("jmp .");
    // }

    if (!(search_for_spc((uint32_t)addr_))) {
        spc_table[spc_count].start_addr = (uint32_t)addr_;
        while(1) {
            test_patch_entry = Restore_Patch_Value((uint32_t)addr_);
            if ((search_for_spc((uint32_t)addr_))) {
                spc_table[spc_count].end_addr = (uint32_t)x86_i_t->start_i;
                spc_count++;
                break;
            } else if (test_patch_entry->patch_addr != 0) {
                if ((test_patch_entry->Value_before_patching != RET) && (test_patch_entry->Value_before_patching != RET_IMM16) && (test_patch_entry->Value_before_patching != CALL_RM32)) {
                    *((uint8_t*)(test_patch_entry->patch_addr)) = test_patch_entry->Value_before_patching;
                } else {
                    put_str("scan_patch error\n");
                    asm("jmp .");
                }
            }
            x86_i_t = x86_it_table;
            x86_i_t->i_addr = (uint8_t*)addr_;
            x86_i_t->start_i = (uint8_t*)addr_;

            x86_dis(x86_i_t);
            addr_ = (void *)x86_i_t->i_addr;
            put_str("\n");
            
            switch (x86_i_t->opcode[0]) {
                case 0x0f:
                switch (x86_i_t->opcode[1]) {
                    case JLE_REL32:
                    spc_table[spc_count].end_addr = (uint32_t)x86_i_t->start_i;
                    spc_count++;
                    state = true;

                    if (x86_i_t->displacement & 0x80000000) {
                        test_addr = (uint32_t)((~((uint32_t)(x86_i_t->displacement))) + 0x1);
                        test_addr = (uint32_t)(x86_i_t->i_addr) - test_addr;
                        if (!(search_for_spc((uint32_t)test_addr))) {
                            patch_table[patch_count].patch_addr = test_addr;
                            patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                            *((uint8_t*)test_addr) = INT3;
                            patch_count++;
                        }
                    } else {
                        test_addr = (uint32_t)(x86_i_t->i_addr) + (uint32_t)(x86_i_t->displacement);
                        if (!(search_for_spc((uint32_t)test_addr))) {
                            patch_table[patch_count].patch_addr = test_addr;
                            patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                            *((uint8_t*)test_addr) = INT3;
                            patch_count++;
                        }
                        
                    }
                    test_addr = (uint32_t)(x86_i_t->i_addr);
                    if (!(search_for_spc((uint32_t)test_addr))) {
                        patch_table[patch_count].patch_addr = test_addr;
                        patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                        *((uint8_t*)test_addr) = INT3;
                        patch_count++;
                    }
                    break;

                    case JNE_REL32:
                    spc_table[spc_count].end_addr = (uint32_t)x86_i_t->start_i;
                    spc_count++;
                    state = true;

                    if (x86_i_t->displacement & 0x80000000) {
                        test_addr = (uint32_t)((~((uint32_t)(x86_i_t->displacement))) + 0x1);
                        test_addr = (uint32_t)(x86_i_t->i_addr) - test_addr;
                        if (!(search_for_spc((uint32_t)test_addr))) {
                            patch_table[patch_count].patch_addr = test_addr;
                            patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                            *((uint8_t*)test_addr) = INT3;
                            patch_count++;
                        }
                    } else {
                        test_addr = (uint32_t)(x86_i_t->i_addr) + (uint32_t)(x86_i_t->displacement);
                        if (!(search_for_spc((uint32_t)test_addr))) {
                            patch_table[patch_count].patch_addr = test_addr;
                            patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                            *((uint8_t*)test_addr) = INT3;
                            patch_count++;
                        }
                        
                    }
                    test_addr = (uint32_t)(x86_i_t->i_addr);
                    if (!(search_for_spc((uint32_t)test_addr))) {
                        patch_table[patch_count].patch_addr = test_addr;
                        patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                        *((uint8_t*)test_addr) = INT3;
                        patch_count++;
                    }
                    break;

                    case JE_REL32:
                    spc_table[spc_count].end_addr = (uint32_t)x86_i_t->start_i;
                    spc_count++;
                    state = true;

                    if (x86_i_t->displacement & 0x80000000) {
                        test_addr = (uint32_t)((~((uint32_t)(x86_i_t->displacement))) + 0x1);
                        test_addr = (uint32_t)(x86_i_t->i_addr) - test_addr;
                        if (!(search_for_spc((uint32_t)test_addr))) {
                            patch_table[patch_count].patch_addr = test_addr;
                            patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                            *((uint8_t*)test_addr) = INT3;
                            patch_count++;
                        }
                    } else {
                        test_addr = (uint32_t)(x86_i_t->i_addr) + (uint32_t)(x86_i_t->displacement);
                        if (!(search_for_spc((uint32_t)test_addr))) {
                            patch_table[patch_count].patch_addr = test_addr;
                            patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                            *((uint8_t*)test_addr) = INT3;
                            patch_count++;
                        }
                        
                    }
                    test_addr = (uint32_t)(x86_i_t->i_addr);
                    if (!(search_for_spc((uint32_t)test_addr))) {
                        patch_table[patch_count].patch_addr = test_addr;
                        patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                        *((uint8_t*)test_addr) = INT3;
                        patch_count++;
                    }
                    break;
                }
                break;

                case CALL_RE32:
                spc_table[spc_count].end_addr = (uint32_t)x86_i_t->start_i;
                spc_count++;
                state = true;
                if (x86_i_t->displacement & 0x80000000) {
                    test_addr = (uint32_t)((~((uint32_t)(x86_i_t->displacement))) + 0x1);
                    test_addr = (uint32_t)(x86_i_t->i_addr) - test_addr;
                    if (!(search_for_spc((uint32_t)test_addr))) {
                        patch_table[patch_count].patch_addr = test_addr;
                        patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                        *((uint8_t*)test_addr) = INT3;
                        patch_count++;
                    }
                } else {
                    test_addr = (uint32_t)(x86_i_t->i_addr) + (uint32_t)(x86_i_t->displacement);
                    if (!(search_for_spc((uint32_t)test_addr))) {
                        patch_table[patch_count].patch_addr = test_addr;
                        patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                        *((uint8_t*)test_addr) = INT3;
                        patch_count++;
                    }
                    
                }
                
                break;

                case IRET_:
                spc_table[spc_count].end_addr = (uint32_t)x86_i_t->start_i;
                spc_count++;
                state = true;
                break;

                case JMP_REL32:
                spc_table[spc_count].end_addr = (uint32_t)x86_i_t->start_i;
                spc_count++;
                state = true;
                if (x86_i_t->displacement & 0x80000000) {
                    test_addr = (uint32_t)((~((uint32_t)(x86_i_t->displacement))) + 0x1);
                    test_addr = (uint32_t)(x86_i_t->i_addr) - test_addr;
                    if (!(search_for_spc((uint32_t)test_addr))) {
                        patch_table[patch_count].patch_addr = test_addr;
                        patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                        *((uint8_t*)test_addr) = INT3;
                        patch_count++;
                    }
                } else {
                    test_addr = (uint32_t)(x86_i_t->i_addr) + (uint32_t)(x86_i_t->displacement);
                    if (!(search_for_spc((uint32_t)test_addr))) {
                        patch_table[patch_count].patch_addr = test_addr;
                        patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                        *((uint8_t*)test_addr) = INT3;
                        patch_count++;
                    }
                    
                }
                break;

                case JE_REL8:
                spc_table[spc_count].end_addr = (uint32_t)x86_i_t->start_i;
                spc_count++;
                state = true;
                if (x86_i_t->displacement & 0x80) {
                    test_addr = (uint8_t)(-(uint8_t)(x86_i_t->displacement));
                    test_addr = (uint32_t)(x86_i_t->i_addr) - test_addr;
                    if (!(search_for_spc((uint32_t)test_addr))) {
                        patch_table[patch_count].patch_addr = test_addr;
                        patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                        *((uint8_t*)test_addr) = INT3;
                        patch_count++;
                    }
                } else {
                    test_addr = (uint32_t)(x86_i_t->i_addr) + (uint32_t)(x86_i_t->displacement);
                    if (!(search_for_spc((uint32_t)test_addr))) {
                        patch_table[patch_count].patch_addr = test_addr;
                        patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                        *((uint8_t*)test_addr) = INT3;
                        patch_count++;
                    }
                }
                test_addr = (uint32_t)(x86_i_t->i_addr);
                if (!(search_for_spc((uint32_t)test_addr))) {
                    patch_table[patch_count].patch_addr = test_addr;
                    patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                    *((uint8_t*)test_addr) = INT3;
                    patch_count++;
                }
                break;

                case JMP_REL8:
                spc_table[spc_count].end_addr = (uint32_t)x86_i_t->start_i;
                spc_count++;
                state = true;
                if (x86_i_t->displacement & 0x80) {
                    test_addr = (uint8_t)(-(uint8_t)(x86_i_t->displacement));
                    test_addr = (uint32_t)(x86_i_t->i_addr) - test_addr;
                    if (!(search_for_spc((uint32_t)test_addr))) {
                        patch_table[patch_count].patch_addr = test_addr;
                        patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                        *((uint8_t*)test_addr) = INT3;
                        patch_count++;
                    }
                } else {
                    test_addr = (uint32_t)(x86_i_t->i_addr) + (uint32_t)(x86_i_t->displacement);
                    if (!(search_for_spc((uint32_t)test_addr))) {
                        patch_table[patch_count].patch_addr = test_addr;
                        patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                        *((uint8_t*)test_addr) = INT3;
                        patch_count++;
                    }
                }
                break;

                case JL_REL8:
                spc_table[spc_count].end_addr = (uint32_t)x86_i_t->start_i;
                spc_count++;
                state = true;
                if (x86_i_t->displacement & 0x80) {
                    test_addr = (uint8_t)(-(uint8_t)(x86_i_t->displacement));
                    test_addr = (uint32_t)(x86_i_t->i_addr) - test_addr;
                    if (!(search_for_spc((uint32_t)test_addr))) {
                        patch_table[patch_count].patch_addr = test_addr;
                        patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                        *((uint8_t*)test_addr) = INT3;
                        patch_count++;
                    }
                } else {
                    test_addr = (uint32_t)(x86_i_t->i_addr) + (uint32_t)(x86_i_t->displacement);
                    if (!(search_for_spc((uint32_t)test_addr))) {
                        patch_table[patch_count].patch_addr = test_addr;
                        patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                        *((uint8_t*)test_addr) = INT3;
                        patch_count++;
                    }
                }
                test_addr = (uint32_t)(x86_i_t->i_addr);
                if (!(search_for_spc((uint32_t)test_addr))) {
                    patch_table[patch_count].patch_addr = test_addr;
                    patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                    *((uint8_t*)test_addr) = INT3;
                    patch_count++;
                }
                break;

                case LOOP_REL8:
                spc_table[spc_count].end_addr = (uint32_t)x86_i_t->start_i;
                spc_count++;
                state = true;

                if (x86_i_t->displacement & 0x80) {
                    test_addr = (uint8_t)(-(uint8_t)(x86_i_t->displacement));
                    put_int(test_addr);
                    put_str("\n");
                    test_addr = (uint32_t)(x86_i_t->i_addr) - test_addr;
                    put_int(test_addr);
                    put_str("\n");
                    //asm("jmp .");
                    if (!(search_for_spc((uint32_t)test_addr))) {
                        put_str("loop______\n");
                        patch_table[patch_count].patch_addr = test_addr;
                        patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                        *((uint8_t*)test_addr) = INT3;
                        patch_count++;
                    }

                } else {
                    put_str("x86_scan error \n");
                    asm("jmp .");
                }
                test_addr = (uint32_t)(x86_i_t->i_addr);
                if (!(search_for_spc((uint32_t)test_addr))) {
                    patch_table[patch_count].patch_addr = test_addr;
                    patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                    *((uint8_t*)test_addr) = INT3;
                    patch_count++;
                }
                break;

                case RET:
                spc_table[spc_count].end_addr = (uint32_t)x86_i_t->start_i;
                spc_count++;
                state = true;

                test_addr = (uint32_t)x86_i_t->start_i;
                patch_table[patch_count].patch_addr = test_addr;
                patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                *((uint8_t*)test_addr) = INT3;
                patch_count++;
                break;

                case JLE_REL8:
                spc_table[spc_count].end_addr = (uint32_t)x86_i_t->start_i;
                spc_count++;
                state = true;
                if (x86_i_t->displacement & 0x80) {
                    test_addr = (uint8_t)(-(uint8_t)(x86_i_t->displacement));
                    test_addr = (uint32_t)(x86_i_t->i_addr) - test_addr;
                    if (!(search_for_spc((uint32_t)test_addr))) {
                        patch_table[patch_count].patch_addr = test_addr;
                        patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                        *((uint8_t*)test_addr) = INT3;
                        patch_count++;
                    }
                } else {
                    test_addr = (uint32_t)(x86_i_t->i_addr) + (uint32_t)(x86_i_t->displacement);
                    if (!(search_for_spc((uint32_t)test_addr))) {
                        patch_table[patch_count].patch_addr = test_addr;
                        patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                        *((uint8_t*)test_addr) = INT3;
                        patch_count++;
                    }
                }
                test_addr = (uint32_t)(x86_i_t->i_addr);
                if (!(search_for_spc((uint32_t)test_addr))) {
                    patch_table[patch_count].patch_addr = test_addr;
                    patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                    *((uint8_t*)test_addr) = INT3;
                    patch_count++;
                }
                break;

                case JNE_REL8:
                spc_table[spc_count].end_addr = (uint32_t)x86_i_t->start_i;
                spc_count++;
                state = true;
                if (x86_i_t->displacement & 0x80) {
                    test_addr = (uint8_t)(-(uint8_t)(x86_i_t->displacement));
                    test_addr = (uint32_t)(x86_i_t->i_addr) - test_addr;
                    if (!(search_for_spc((uint32_t)test_addr))) {
                        patch_table[patch_count].patch_addr = test_addr;
                        patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                        *((uint8_t*)test_addr) = INT3;
                        patch_count++;
                    }
                } else {
                    test_addr = (uint32_t)(x86_i_t->i_addr) + (uint32_t)(x86_i_t->displacement);
                    if (!(search_for_spc((uint32_t)test_addr))) {
                        patch_table[patch_count].patch_addr = test_addr;
                        patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                        *((uint8_t*)test_addr) = INT3;
                        patch_count++;
                    }
                }
                test_addr = (uint32_t)(x86_i_t->i_addr);
                if (!(search_for_spc((uint32_t)test_addr))) {
                    patch_table[patch_count].patch_addr = test_addr;
                    patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                    *((uint8_t*)test_addr) = INT3;
                    patch_count++;
                }
                break;

                case JBE_REL8:
                spc_table[spc_count].end_addr = (uint32_t)x86_i_t->start_i;
                spc_count++;
                state = true;
                if (x86_i_t->displacement & 0x80) {
                    test_addr = (uint8_t)(-(uint8_t)(x86_i_t->displacement));
                    test_addr = (uint32_t)(x86_i_t->i_addr) - test_addr;
                    if (!(search_for_spc((uint32_t)test_addr))) {
                        patch_table[patch_count].patch_addr = test_addr;
                        patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                        *((uint8_t*)test_addr) = INT3;
                        patch_count++;
                    }
                } else {
                    test_addr = (uint32_t)(x86_i_t->i_addr) + (uint32_t)(x86_i_t->displacement);
                    if (!(search_for_spc((uint32_t)test_addr))) {
                        patch_table[patch_count].patch_addr = test_addr;
                        patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                        *((uint8_t*)test_addr) = INT3;
                        patch_count++;
                    }
                }
                test_addr = (uint32_t)(x86_i_t->i_addr);
                if (!(search_for_spc((uint32_t)test_addr))) {
                    patch_table[patch_count].patch_addr = test_addr;
                    patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                    *((uint8_t*)test_addr) = INT3;
                    patch_count++;
                }
                break;

                case RET_IMM16:
                spc_table[spc_count].end_addr = (uint32_t)x86_i_t->start_i;
                spc_count++;
                state = true;

                test_addr = (uint32_t)x86_i_t->start_i;
                patch_table[patch_count].patch_addr = test_addr;
                patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                *((uint8_t*)test_addr) = INT3;
                patch_count++;

                break;

                case CALL_RM32:
                spc_table[spc_count].end_addr = (uint32_t)x86_i_t->start_i;
                spc_count++;
                state = true;

                test_addr = (uint32_t)x86_i_t->start_i;
                patch_table[patch_count].patch_addr = test_addr;
                patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                *((uint8_t*)test_addr) = INT3;
                patch_count++;

                break;

                case JA_REL8:
                spc_table[spc_count].end_addr = (uint32_t)x86_i_t->start_i;
                spc_count++;
                state = true;

                if (x86_i_t->displacement & 0x80) {
                    test_addr = (uint8_t)(-(uint8_t)(x86_i_t->displacement));
                    test_addr = (uint32_t)(x86_i_t->i_addr) - test_addr;
                    if (!(search_for_spc((uint32_t)test_addr))) {
                        patch_table[patch_count].patch_addr = test_addr;
                        patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                        *((uint8_t*)test_addr) = INT3;
                        patch_count++;
                    }
                } else {
                    test_addr = (uint32_t)(x86_i_t->i_addr) + (uint32_t)(x86_i_t->displacement);
                    if (!(search_for_spc((uint32_t)test_addr))) {
                        patch_table[patch_count].patch_addr = test_addr;
                        patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                        *((uint8_t*)test_addr) = INT3;
                        patch_count++;
                    }
                }
                test_addr = (uint32_t)(x86_i_t->i_addr);
                if (!(search_for_spc((uint32_t)test_addr))) {
                    patch_table[patch_count].patch_addr = test_addr;
                    patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                    *((uint8_t*)test_addr) = INT3;
                    patch_count++;
                }

                break;

                case JAE_REL8:
                spc_table[spc_count].end_addr = (uint32_t)x86_i_t->start_i;
                spc_count++;
                state = true;

                if (x86_i_t->displacement & 0x80) {
                    test_addr = (uint8_t)(-(uint8_t)(x86_i_t->displacement));
                    test_addr = (uint32_t)(x86_i_t->i_addr) - test_addr;
                    if (!(search_for_spc((uint32_t)test_addr))) {
                        patch_table[patch_count].patch_addr = test_addr;
                        patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                        *((uint8_t*)test_addr) = INT3;
                        patch_count++;
                    }
                } else {
                    test_addr = (uint32_t)(x86_i_t->i_addr) + (uint32_t)(x86_i_t->displacement);
                    if (!(search_for_spc((uint32_t)test_addr))) {
                        patch_table[patch_count].patch_addr = test_addr;
                        patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                        *((uint8_t*)test_addr) = INT3;
                        patch_count++;
                    }
                }
                test_addr = (uint32_t)(x86_i_t->i_addr);
                if (!(search_for_spc((uint32_t)test_addr))) {
                    patch_table[patch_count].patch_addr = test_addr;
                    patch_table[patch_count].Value_before_patching = *((uint8_t*)test_addr);
                    *((uint8_t*)test_addr) = INT3;
                    patch_count++;
                }

                break;
            }
            x86_it_table++;

            if (state) {
                break;;
            }
        }
    }
    
}

void x86_scan_patch_init(void) {
    Restore_Patch_Value_test_patch_entry.patch_addr = 0x0;
    Restore_Patch_Value_test_patch_entry.Value_before_patching = 0x0;
    x86_it_table = get_kernel_pages(400);
    put_str("x86 scan patch init\n");
}